package Protocol;

import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

//CIPHER / GENERATORS
import javax.crypto.SecretKey;

//KEY SPECIFICATIONS

import javax.crypto.spec.IvParameterSpec;

//EXCEPTIONS
import java.security.InvalidAlgorithmParameterException;
import java.io.UnsupportedEncodingException;

public class CommunicationProtocol {

	Cipher ecipher;
	Cipher dcipher;

	/**
	 * Constructor used to create this object. Responsible for setting and
	 * initializing this object's encrypter and decrypter Chipher instances
	 * given a Secret Key and algorithm.
	 * 
	 * @param key
	 *            Secret Key used to initialize both the encrypter and decrypter
	 *            instances.
	 * @param algorithm
	 *            Which algorithm to use for creating the encrypter and
	 *            decrypter instances.
	 * @throws InvalidAlgorithmParameterException
	 */
	public CommunicationProtocol(SecretKey key)
			throws InvalidAlgorithmParameterException {
		try {
			// Obtain a cipher engine that can perform encryption and decryption
			// by implementing the AES/CBC/PKCS5Padding algorithm.
			ecipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
			dcipher = Cipher.getInstance("AES/CBC/PKCS5Padding");

			// Use SecureRandom to get a random IvParameterSpec
			int blksz = ecipher.getBlockSize();
			byte[] ivBytes = new byte[blksz];
			SecureRandom random = new SecureRandom();
			random.nextBytes(ivBytes);
			IvParameterSpec ivSpec = new IvParameterSpec(ivBytes);

			// Initialize the cipher to encrypt or decrypt data.
			ecipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);
			dcipher.init(Cipher.DECRYPT_MODE, key, ivSpec);

		} catch (NoSuchPaddingException e) {
			System.out.println("EXCEPTION: NoSuchPaddingException");
		} catch (NoSuchAlgorithmException e) {
			System.out.println("EXCEPTION: NoSuchAlgorithmException");
		} catch (InvalidKeyException e) {
			System.out.println("EXCEPTION: InvalidKeyException");
		}
	}

	public EncryptMessageType encryptMessage(String sendMessage)
			throws UnsupportedEncodingException, IllegalBlockSizeException,
			BadPaddingException, NoSuchAlgorithmException {
		EncryptMessageType encryptMsg = new EncryptMessageType();

		// Covert the sendMessage from String to byte Array
		byte[] utf8 = sendMessage.getBytes("UTF8");

		// Encrypt or decrypt the data in the input array as well as any data
		// that has been previously buffered in the cipher engine.
		byte[] encSendMessage = ecipher.doFinal(utf8);

		encryptMsg.setEncryMsg(encSendMessage);

		// hash the sendMessage
		byte[] hashedMsg = this.hashMessage(sendMessage);

		encryptMsg.setHashedMsg(hashedMsg);

		return encryptMsg;
	}

	public String decryptMessage(EncryptMessageType receivedMessage)
			throws UnsupportedEncodingException, IllegalBlockSizeException,
			BadPaddingException, NoSuchAlgorithmException {

		byte[] encryMsg = receivedMessage.getEncryMsg();
		byte[] hashedMsg = receivedMessage.getHashedMsg();

		// Decrypt
		byte[] utf8 = dcipher.doFinal(encryMsg);

		// Decode using utf-8
		String decrymsg = new String(utf8, "UTF8");

		// compute the hash of the decrypted message
		byte[] computeHashedMsg = this.hashMessage(decrymsg);

		// compare the hash of receivedMsg to the hash in the message to check
		// the intergrity of the message
		String stringHashedMsg = new String(hashedMsg, "UTF8");
		String stringComputeHashedMsg = new String(computeHashedMsg, "UTF8");

		if (stringHashedMsg.equals(stringComputeHashedMsg)) {
			return decrymsg;
		} else {
			return null;
		}
	}

	// This function is used to compute the hash of message
	public byte[] hashMessage(String MessageToHash)
			throws NoSuchAlgorithmException, UnsupportedEncodingException {

		MessageDigest messageDigest = MessageDigest.getInstance("SHA-512");
		messageDigest.update(MessageToHash.getBytes("UTF8"));
		byte[] MessageHashed = messageDigest.digest();
		//
		return MessageHashed;
	}

}